import argparse
import json
import os
import numpy as np
import cv2
from PIL import Image
from abc import ABC, abstractmethod
from skimage.metrics import peak_signal_noise_ratio, structural_similarity
import time

class DenoiseUnit(ABC):
    """算法单元基类"""
    
    def __init__(self, name="DenoiseUnit"):
        self.name = name
        self.params = {}
    
    @abstractmethod
    def process(self, frame, **kwargs):
        """处理单帧图像"""
        pass
    
    def set_params(self, **kwargs):
        """设置算法参数"""
        for key, value in kwargs.items():
            if key in self.params:
                self.params[key] = value
            else:
                print(f"警告: 参数 '{key}' 不在 {self.name} 的参数列表中")
    
    def get_params(self):
        """获取当前参数"""
        return self.params.copy()
    
    def to_dict(self):
        """转换为字典格式，用于JSON序列化"""
        return {
            "name": self.name,
            "params": self.params
        }


class TemporalAverage(DenoiseUnit):
    """时域平均降噪"""
    
    def __init__(self, buffer_size=5):
        super().__init__("TemporalAverage")
        self.params = {"buffer_size": buffer_size}
        self.frame_buffer = []
    
    def process(self, frame, **kwargs):
        if not isinstance(frame, np.ndarray):
            frame = np.array(frame)
        
        # 添加到帧缓冲区
        self.frame_buffer.append(frame.copy())
        if len(self.frame_buffer) > self.params["buffer_size"]:
            self.frame_buffer.pop(0)
        
        # 计算平均值
        if len(self.frame_buffer) > 0:
            result = np.mean(self.frame_buffer, axis=0).astype(np.uint8)
            return result
        return frame


class TemporalMedian(DenoiseUnit):
    """时域中值滤波"""
    
    def __init__(self, buffer_size=5):
        super().__init__("TemporalMedian")
        self.params = {"buffer_size": buffer_size}
        self.frame_buffer = []
    
    def process(self, frame, **kwargs):
        if not isinstance(frame, np.ndarray):
            frame = np.array(frame)
        
        # 添加到帧缓冲区
        self.frame_buffer.append(frame.copy())
        if len(self.frame_buffer) > self.params["buffer_size"]:
            self.frame_buffer.pop(0)
        
        # 计算中值
        if len(self.frame_buffer) > 0:
            result = np.median(self.frame_buffer, axis=0).astype(np.uint8)
            return result
        return frame


class RecursiveFilter(DenoiseUnit):
    """递归滤波"""
    
    def __init__(self, alpha=0.8):
        super().__init__("RecursiveFilter")
        self.params = {"alpha": alpha}
        self.previous_frame = None
    
    def process(self, frame, **kwargs):
        if not isinstance(frame, np.ndarray):
            frame = np.array(frame)
        
        if self.previous_frame is None:
            self.previous_frame = frame.astype(np.float32)
            return frame
        
        # 递归滤波: I_out = α * I_prev + (1-α) * I_current
        alpha = self.params["alpha"]
        result = alpha * self.previous_frame + (1 - alpha) * frame.astype(np.float32)
        result = np.clip(result, 0, 255).astype(np.uint8)
        self.previous_frame = result.astype(np.float32)
        
        return result


class GuidedFilter(DenoiseUnit):
    """引导滤波"""
    
    def __init__(self, radius=5, eps=0.01):
        super().__init__("GuidedFilter")
        self.params = {"radius": radius, "eps": eps}
    
    def process(self, frame, **kwargs):
        if not isinstance(frame, np.ndarray):
            frame = np.array(frame)
        
        # 如果是彩色图像，分别处理每个通道
        if len(frame.shape) == 3:
            result = np.zeros_like(frame)
            for i in range(frame.shape[2]):
                result[:, :, i] = cv2.ximgproc.guidedFilter(
                    frame[:, :, i], frame[:, :, i], 
                    self.params["radius"], self.params["eps"]
                )
            return result
        else:
            return cv2.ximgproc.guidedFilter(
                frame, frame, 
                self.params["radius"], self.params["eps"]
            )


class BilateralFilter(DenoiseUnit):
    """双边滤波"""
    
    def __init__(self, d=5, sigma_color=75, sigma_space=75):
        super().__init__("BilateralFilter")
        self.params = {
            "d": d, 
            "sigma_color": sigma_color, 
            "sigma_space": sigma_space
        }
    
    def process(self, frame, **kwargs):
        if not isinstance(frame, np.ndarray):
            frame = np.array(frame)
        
        return cv2.bilateralFilter(
            frame, 
            self.params["d"], 
            self.params["sigma_color"], 
            self.params["sigma_space"]
        )


class NonLocalMeans(DenoiseUnit):
    """非局部均值去噪"""
    
    def __init__(self, h=10, template_window_size=7, search_window_size=21):
        super().__init__("NonLocalMeans")
        self.params = {
            "h": h,
            "template_window_size": template_window_size,
            "search_window_size": search_window_size
        }
    
    def process(self, frame, **kwargs):
        if not isinstance(frame, np.ndarray):
            frame = np.array(frame)
        
        return cv2.fastNlMeansDenoising(
            frame,
            h=self.params["h"],
            templateWindowSize=self.params["template_window_size"],
            searchWindowSize=self.params["search_window_size"]
        )


class DenoisePipeline:
    """降噪处理管道"""
    
    def __init__(self):
        self.units = []
    
    def add_unit(self, unit):
        """添加处理单元"""
        self.units.append(unit)
        return self
    
    def remove_unit(self, index):
        """移除处理单元"""
        if 0 <= index < len(self.units):
            self.units.pop(index)
    
    def process(self, frame):
        """处理单帧图像"""
        result = frame
        for unit in self.units:
            result = unit.process(result)
        return result
    
    def process_sequence(self, frames):
        """处理帧序列"""
        results = []
        for frame in frames:
            results.append(self.process(frame))
        return results
    
    def get_config(self):
        """获取当前配置"""
        config = []
        for unit in self.units:
            config.append(unit.to_dict())
        return config
    
    def save_config(self, filepath):
        """保存配置到JSON文件"""
        config = self.get_config()
        with open(filepath, 'w') as f:
            json.dump(config, f, indent=4)
    
    @classmethod
    def from_config(cls, config):
        """从配置创建管道"""
        pipeline = cls()
        unit_map = {
            "TemporalAverage": TemporalAverage,
            "TemporalMedian": TemporalMedian,
            "RecursiveFilter": RecursiveFilter,
            "GuidedFilter": GuidedFilter,
            "BilateralFilter": BilateralFilter,
            "NonLocalMeans": NonLocalMeans
        }
        
        for unit_config in config:
            unit_name = unit_config["name"]
            if unit_name in unit_map:
                unit = unit_map[unit_name]()
                unit.set_params(**unit_config["params"])
                pipeline.add_unit(unit)
            else:
                print(f"警告: 未知算法单元 '{unit_name}'，已跳过")
        
        return pipeline


def load_image_sequence(images_path, file_patterns, limit=None):
    """加载图像序列"""
    import glob
    
    # 获取所有匹配的图像文件
    image_files = []
    for pat in file_patterns:
        image_files += glob.glob(f"{images_path}/{pat}")
    image_files = sorted(image_files)
    
    if limit:
        image_files = image_files[:limit]
    
    # 加载图像
    images = []
    for file_path in image_files:
        img = Image.open(file_path)
        images.append(img)
    
    return images, image_files


def calculate_metrics(original, denoised):
    """计算图像质量指标"""
    if not isinstance(original, np.ndarray):
        original = np.array(original)
    if not isinstance(denoised, np.ndarray):
        denoised = np.array(denoised)
    
    # 转换为灰度图像计算SSIM
    if len(original.shape) == 3:
        original_gray = cv2.cvtColor(original, cv2.COLOR_RGB2GRAY)
        denoised_gray = cv2.cvtColor(denoised, cv2.COLOR_RGB2GRAY)
    else:
        original_gray = original
        denoised_gray = denoised
    
    # 计算PSNR
    psnr = peak_signal_noise_ratio(original, denoised)
    
    # 计算SSIM
    ssim = structural_similarity(original_gray, denoised_gray)
    
    return {"PSNR": psnr, "SSIM": ssim}

def usage():
    help_epilog = """
算法与参数说明
- TemporalAverage（时域平均）:
  参数:
    buffer_size: int, 缓冲区帧数(>0)，对最近 buffer_size 帧取像素均值，默认 5。
  适用: 稳定场景的随机噪声抑制；会引入运动模糊。

- TemporalMedian（时域中值滤波）:
  参数:
    buffer_size: int, 缓冲区帧数(>0)，对最近 buffer_size 帧取像素中值，默认 5。
  适用: 脉冲噪声/散斑噪声；相对平均法更能保边；运动大时仍会模糊。

- RecursiveFilter（递归滤波/指数滑动平均）:
  参数:
    alpha: float, 0~1，历史权重，越大越平滑，响应越慢，默认 0.8。
  适用: 在线处理/低延迟场景；在运动时的拖影较可控。

- GuidedFilter（引导滤波）:
  参数:
    radius: int, 窗半径(>0)，默认 5。
    eps: float, 正则项(>0)，默认 0.01。
  适用: 保边平滑；依赖 OpenCV ximgproc 模块（opencv-contrib-python）。
  注意: 彩色图像按通道独立处理。

- BilateralFilter（双边滤波）:
  参数:
    d: int, 邻域直径，默认 5。
    sigma_color: float, 颜色空间标准差，默认 75。
    sigma_space: float, 坐标空间标准差，默认 75。
  适用: 保边去噪；参数过大易过度平滑。

- NonLocalMeans（非局部均值去噪）:
  参数:
    h: float, 过滤强度，默认 10。
    template_window_size: int, 模板窗口大小(奇数)，默认 7。
    search_window_size: int, 搜索窗口大小(奇数)，默认 21。
  适用: 纹理保留较好；计算量较大；使用 cv2.fastNlMeansDenoising（灰度）。

算法配置文件使用方法
- 配置文件为 JSON 数组，每个元素为一个处理单元，按先后顺序依次执行。
- 每个单元包含:
  name: 字符串，对应上述算法名称之一
  params: 对象，包含该算法可用的参数键值

示例 config.json:
[
  {
    "name": "TemporalMedian",
    "params": { "buffer_size": 5 }
  },
  {
    "name": "RecursiveFilter",
    "params": { "alpha": 0.7 }
  },
  {
    "name": "BilateralFilter",
    "params": { "d": 5, "sigma_color": 60, "sigma_space": 60 }
  }
]

使用示例
- 基本运行:
  python script.py -i ./images -o ./out -c ./config.json
- 限制只处理前 20 张并计算指标:
  python script.py -i ./images -o ./out -c ./config.json -l 20 -m
- 指定匹配模式(逗号分隔多个，或在命令中传多次 -p):
  python script.py -i ./images -o ./out -c ./config.json -p *.png -p *.jpg

注意事项
- GuidedFilter 需要安装 opencv-contrib-python。
- NonLocalMeans 当前对灰度图生效；若输入为彩色，可先转灰度或扩展为彩色版本。
- 时域方法依赖“时间缓冲”；单张图片运行时，效果等同于原图或单次滤波。
"""
    print(help_epilog)

def main():
    parser = argparse.ArgumentParser(description="多帧时域图片降噪算法")
    parser.add_argument("--images", "-i", required=True, help="输入图片序列路径")
    parser.add_argument("--output", "-o", required=True, help="降噪输出路径")
    parser.add_argument("--config", "-c", required=True, help="算法配置JSON文件")
    parser.add_argument("--file_patterns", "-p", default=["*.jpg","*.png","*.bmp"], help="文件匹配模式")
    parser.add_argument("--limit", "-l", type=int, default=None, help="处理图像数量限制")
    parser.add_argument("--metrics", "-m", action="store_true", help="是否计算质量指标")
    parser.add_argument("--usage", "-u", action="store_true", help="是否显示用法")
    
    args = parser.parse_args()
    if args.usage:
        usage()
        return
    # 创建输出目录
    os.makedirs(args.output, exist_ok=True)
    
    # 加载配置
    with open(args.config, 'r') as f:
        config = json.load(f)
    
    # 创建处理管道
    pipeline = DenoisePipeline.from_config(config)
    
    # 加载图像序列
    print("加载图像序列...")
    images, image_files = load_image_sequence(args.images, args.file_patterns, args.limit)
    
    # 处理图像序列
    print("处理图像序列...")
    start_time = time.time()
    results = pipeline.process_sequence(images)
    processing_time = time.time() - start_time
    
    # 保存结果图像
    print("保存结果图像...")
    for i, (result, original_path) in enumerate(zip(results, image_files)):
        filename = os.path.basename(original_path)
        output_path = os.path.join(args.output, f"denoised_{filename}")
        
        if isinstance(result, np.ndarray):
            result_img = Image.fromarray(result)
        else:
            result_img = result
            
        result_img.save(output_path)
    
    # 计算并保存质量指标
    metrics_results = []
    if args.metrics:
        print("计算质量指标...")
        for i, (original, denoised) in enumerate(zip(images, results)):
            metrics = calculate_metrics(original, denoised)
            metrics_results.append({
                "image": os.path.basename(image_files[i]),
                "metrics": metrics
            })
    
    # 保存处理结果到文本文件
    result_file = os.path.join(args.output, "result.txt")
    with open(result_file, 'w') as f:
        f.write("多帧时域图片降噪处理结果\n")
        f.write("=" * 50 + "\n\n")
        
        f.write(f"输入路径: {args.images}\n")
        f.write(f"输出路径: {args.output}\n")
        f.write(f"配置文件: {args.config}\n")
        f.write(f"处理图像数量: {len(images)}\n")
        f.write(f"处理时间: {processing_time:.2f} 秒\n")
        f.write(f"平均每帧处理时间: {processing_time/len(images):.4f} 秒\n\n")
        
        f.write("算法配置:\n")
        for i, unit in enumerate(pipeline.units):
            f.write(f"{i+1}. {unit.name}: {unit.get_params()}\n")
        f.write("\n")
        
        if args.metrics:
            f.write("质量指标:\n")
            for result in metrics_results:
                f.write(f"{result['image']}:\n")
                f.write(f"  PSNR: {result['metrics']['PSNR']:.4f}\n")
                f.write(f"  SSIM: {result['metrics']['SSIM']:.4f}\n")
            
            # 计算平均指标
            avg_psnr = np.mean([r['metrics']['PSNR'] for r in metrics_results])
            avg_ssim = np.mean([r['metrics']['SSIM'] for r in metrics_results])
            f.write(f"\n平均指标:\n")
            f.write(f"  PSNR: {avg_psnr:.4f}\n")
            f.write(f"  SSIM: {avg_ssim:.4f}\n")
    
    print(f"处理完成! 结果已保存到 {args.output}")
    print(f"详细信息请查看 {result_file}")


if __name__ == "__main__":
    main()